package com.skivingcloud.ai.service.impl;

import com.alibaba.cloud.ai.tongyi.audio.api.SpeechClient;
import com.skivingcloud.ai.service.AiCommonApiService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImageClient;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.nio.ByteBuffer;
import java.util.Map;

/**
 * @author hushouquan
 * @since 2024/6/27
 */
@Service("tongYiAiApiService")
public class TongYiAiApiServiceImpl implements AiCommonApiService{
    private static final Logger log = LoggerFactory.getLogger(TongYiAiApiServiceImpl.class);
    
    private final ChatClient chatClient;
    private final ImageClient imageClient;
    private final SpeechClient speechClient;
    private final StreamingChatClient streamingChatClient;
    
    @Autowired
    public TongYiAiApiServiceImpl(ChatClient chatClient, ImageClient imageClient, SpeechClient speechClient, StreamingChatClient streamingChatClient) {
        this.chatClient = chatClient;
        this.imageClient = imageClient;
        this.speechClient = speechClient;
        this.streamingChatClient = streamingChatClient;
    }
    /**
     * 基本问答
     *
     * @param message 参数
     */
    @Override
    public String completion(String message){
        Prompt prompt = new Prompt(new UserMessage(message));
        return chatClient.call(prompt).getResult().getOutput().getContent();
    }
    
    /**
     * 文生图
     *
     * @param imgPrompt 参数
     */
    @Override
    public ImageResponse genImg(String imgPrompt){
        ImagePrompt prompt = new ImagePrompt(imgPrompt);
        return imageClient.call(prompt);
    }
    
    /**
     * 语音合成
     *
     * @param text 参数
     */
    @Override
    public ByteBuffer genAudio(String text){
        // save的代码省略，就是将音频保存到本地而已
        // return save(resWAV, SpeechSynthesisAudioFormat.WAV.getValue());
        return speechClient.call(text);
    }
    
    /**
     * 连续问答
     *
     * @param message 参数
     * @return 结果
     */
    @Override
    public Map<String, String> streamCompletion(String message){
        StringBuilder fullContent = new StringBuilder();
        
        streamingChatClient.stream(new Prompt(message))
                .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))
                .map(content -> content.getOutput().getContent())
                .doOnNext(fullContent::append)
                .last()
                .map(lastContent -> Map.of(message, fullContent.toString()))
                .block();
        
        log.info(fullContent.toString());
        
        return Map.of(message, fullContent.toString());
    }
}
